from base import base
import os
from utils import Utils
import pandas as pd
from sklearn.metrics import roc_auc_score
class trainLR(base):
    def __init__(self, cfg) -> None:
        super().__init__(cfg)

    def train_one_epoch(self, model, loader):
        model(loader())

    def valid_one_epoch(self, model, loader):
        result = model.test(loader())
        if not os.path.exists(os.path.join(self.train_cfg.get('result_path'), Utils.current_time())):
            os.makedirs(os.path.join(self.train_cfg.get('result_path'), Utils.current_time()), mode=0o755)
        save_path = os.path.join(self.train_cfg.get('result_path'), Utils.current_time(), Utils.current_time_second() + 'submission.csv')
        result.to_csv(save_path, index=None)
        return save_path
    
    def cal_auc(self, groud_truth_path, predict_path):
        label = pd.read_csv(groud_truth_path, sep=',')
        predict = pd.read_csv(predict_path, sep=',')
        data = label.merge(predict, on='index', how='left')
        score = roc_auc_score(data['like'], data['predict'])
        print("auc score: ", score)

    